use crate::proxy::ServiceProxy;
use async_trait::async_trait;
use rustcloud_core::{
    CircuitBreaker, LoadBalancer, Request, Response, ServiceError, ServiceInstance, 
    ServiceRegistry, ServiceResult, Transport, TransportConfig, CircuitBreakerState,
    ServiceWatcher, CircuitBreakerMetrics
};
use rustcloud_circuit_breaker::{DefaultCircuitBreaker, CircuitBreakerConfig};
use std::sync::Arc;
use std::time::Duration;
use std::sync::Mutex;
use std::collections::HashMap;

// --- Enhanced Mocks ---
struct MockRegistry {
    instances: Arc<Mutex<HashMap<String, Vec<ServiceInstance>>>>,
}

impl MockRegistry {
    fn new() -> Self {
        Self {
            instances: Arc::new(Mutex::new(HashMap::new())),
        }
    }
    
    fn add_instance(&self, service_name: &str, instance: ServiceInstance) {
        let mut instances = self.instances.lock().unwrap();
        instances.entry(service_name.to_string()).or_insert_with(Vec::new).push(instance);
    }
    
    fn clear_instances(&self) {
        self.instances.lock().unwrap().clear();
    }
}

#[async_trait]
impl ServiceRegistry for MockRegistry {
    async fn register(&self, _instance: ServiceInstance) -> ServiceResult<()> { Ok(()) }
    async fn deregister(&self, _instance: ServiceInstance) -> ServiceResult<()> { Ok(()) }
    async fn discover(&self, service_name: &str) -> ServiceResult<Vec<ServiceInstance>> {
        let instances = self.instances.lock().unwrap();
        Ok(instances.get(service_name).cloned().unwrap_or_default())
    }
    async fn health_check(&self, service_name: &str) -> ServiceResult<Vec<ServiceInstance>> {
        let instances = self.discover(service_name).await?;
        Ok(instances.into_iter().filter(|i| i.healthy).collect())
    }
    async fn watch(&self, _service_name: &str) -> ServiceResult<Box<dyn ServiceWatcher + Send>> { 
        unimplemented!("Mock registry doesn't support watching") 
    }
}

struct MockLoadBalancer {
    selection_index: Arc<Mutex<usize>>,
}

impl MockLoadBalancer {
    fn new() -> Self {
        Self {
            selection_index: Arc::new(Mutex::new(0)),
        }
    }
}

#[async_trait]
impl LoadBalancer for MockLoadBalancer {
    async fn select(&self, _service_name: &str, instances: &[ServiceInstance]) -> ServiceResult<Option<ServiceInstance>> {
        if instances.is_empty() {
            return Ok(None);
        }
        let mut index = self.selection_index.lock().unwrap();
        let selected = instances.get(*index).cloned();
        *index = (*index + 1) % instances.len();
        Ok(selected)
    }
    
    fn name(&self) -> &str { "mock" }
    
    fn get_strategy_info(&self) -> rustcloud_core::StrategyInfo {
        rustcloud_core::StrategyInfo {
            name: "Mock Load Balancer".to_string(),
            description: "Mock load balancer for testing".to_string(),
            requires_session_affinity: false,
            supports_weighted_routing: false,
        }
    }
}

struct MockTransport {
    responses: Arc<Mutex<HashMap<String, Response>>>,
    call_count: Arc<Mutex<usize>>,
    should_fail: Arc<Mutex<bool>>,
}

impl MockTransport {
    fn new() -> Self {
        Self {
            responses: Arc::new(Mutex::new(HashMap::new())),
            call_count: Arc::new(Mutex::new(0)),
            should_fail: Arc::new(Mutex::new(false)),
        }
    }
    
    fn set_response(&self, path: &str, response: Response) {
        self.responses.lock().unwrap().insert(path.to_string(), response);
    }
    
    fn set_should_fail(&self, should_fail: bool) {
        *self.should_fail.lock().unwrap() = should_fail;
    }
    
    fn get_call_count(&self) -> usize {
        *self.call_count.lock().unwrap()
    }
    
    fn reset_call_count(&self) {
        *self.call_count.lock().unwrap() = 0;
    }
}

#[async_trait]
impl Transport for MockTransport {
    async fn send(&self, request: Request, _instance: &ServiceInstance) -> ServiceResult<Response> {
        *self.call_count.lock().unwrap() += 1;
        
        if *self.should_fail.lock().unwrap() {
            return Err(ServiceError::NetworkError("Simulated network failure".to_string()));
        }
        
        let responses = self.responses.lock().unwrap();
        if let Some(response) = responses.get(&request.path) {
            Ok(response.clone())
        } else {
            Ok(Response::new(200))
        }
    }
    
    async fn send_batch(&self, requests: Vec<(Request, &ServiceInstance)>) -> Vec<ServiceResult<Response>> {
        let mut results = Vec::new();
        for (request, instance) in requests {
            results.push(self.send(request, instance).await);
        }
        results
    }
    
    fn name(&self) -> &str { "mock" }
}

struct FailingMockTransport;

#[async_trait]
impl Transport for FailingMockTransport {
    async fn send(&self, _request: Request, _instance: &ServiceInstance) -> ServiceResult<Response> {
        Err(ServiceError::NetworkError("Simulated network failure".to_string()))
    }
    async fn send_batch(&self, _requests: Vec<(Request, &ServiceInstance)>) -> Vec<ServiceResult<Response>> { 
        unimplemented!() 
    }
    fn name(&self) -> &str { "failing_mock_transport" }
}

// --- Test Setup ---
fn setup_proxy(transport: Arc<dyn Transport>, cb_config: CircuitBreakerConfig) -> ServiceProxy<DefaultCircuitBreaker> {
    let registry = Arc::new(MockRegistry::new());
    let load_balancer = Arc::new(MockLoadBalancer::new());
    let circuit_breaker = Arc::new(DefaultCircuitBreaker::new(cb_config));
    ServiceProxy::new(registry, load_balancer, transport, circuit_breaker)
}

fn setup_proxy_with_registry(registry: Arc<MockRegistry>, transport: Arc<dyn Transport>) -> ServiceProxy<DefaultCircuitBreaker> {
    let load_balancer = Arc::new(MockLoadBalancer::new());
    let circuit_breaker = Arc::new(DefaultCircuitBreaker::new(CircuitBreakerConfig::default()));
    ServiceProxy::new(registry, load_balancer, transport, circuit_breaker)
}

#[tokio::test]
async fn test_service_proxy_call_success() {
    let transport = Arc::new(MockTransport::new());
    let mut success_response = Response::new(200);
    success_response.set_body(b"success".to_vec());
    transport.set_response("/hello", success_response);
    
    let registry = Arc::new(MockRegistry::new());
    registry.add_instance("test-service", ServiceInstance::new("test-service", "localhost", 8080, "http"));
    
    let proxy = setup_proxy_with_registry(registry, transport.clone());
    let config = TransportConfig::default();
    let request = Request::new("test-service", "GET", "/hello");
    
    let response = proxy.call(request, &config).await.unwrap();
    assert_eq!(response.status, 200);
    assert_eq!(response.body, Some(b"success".to_vec()));
    assert_eq!(transport.get_call_count(), 1);
}

#[tokio::test]
async fn test_service_proxy_no_available_instances() {
    let transport = Arc::new(MockTransport::new());
    let registry = Arc::new(MockRegistry::new());
    // 不添加任何实侌
    
    let proxy = setup_proxy_with_registry(registry, transport.clone());
    let config = TransportConfig::default();
    let request = Request::new("nonexistent-service", "GET", "/hello");
    
    let result = proxy.call(request, &config).await;
    assert!(result.is_err());
    
    match result.unwrap_err() {
        ServiceError::NoAvailableInstances(service) => {
            assert_eq!(service, "nonexistent-service");
        }
        _ => panic!("Expected NoAvailableInstances error"),
    }
    
    assert_eq!(transport.get_call_count(), 0); // 没有实侌时不应该调用传输层
}

#[tokio::test]
async fn test_circuit_breaker_opens_and_rejects_calls() {
    let cb_config = CircuitBreakerConfig {
        failure_threshold: 2,
        recovery_timeout: Duration::from_millis(10000), // 10秒
        ..Default::default()
    };
    
    let registry = Arc::new(MockRegistry::new());
    registry.add_instance("test-service", ServiceInstance::new("test-service", "localhost", 8080, "http"));
    
    let proxy = setup_proxy(Arc::new(FailingMockTransport), cb_config);
    let transport_config = TransportConfig { retries: 1, ..Default::default() }; // 总共2次尝试
    let request = Request::new("test-service", "GET", "/hello");

    // 第一次调用失败，第二次尝试也失败，触发熔断器
    let result = proxy.call(request.clone(), &transport_config).await;
    assert!(result.is_err());

    // 立即再次调用，应该被熔断器拒绝
    let result2 = proxy.call(request.clone(), &transport_config).await;
    assert!(result2.is_err());
    match result2.err().unwrap() {
        ServiceError::CircuitBreakerOpen(_) => { /* Success */ }
        e => panic!("Expected CircuitBreakerOpen error, got {:?}", e),
    }
}

#[tokio::test]
async fn test_service_proxy_retry_logic() {
    let transport = Arc::new(MockTransport::new());
    transport.set_should_fail(true); // 设置传输失败
    
    let registry = Arc::new(MockRegistry::new());
    registry.add_instance("test-service", ServiceInstance::new("test-service", "localhost", 8080, "http"));
    
    let proxy = setup_proxy_with_registry(registry, transport.clone());
    let config = TransportConfig {
        retries: 2,
        retry_delay: 10,
        ..Default::default()
    };
    let request = Request::new("test-service", "GET", "/hello");
    
    let result = proxy.call(request, &config).await;
    assert!(result.is_err());
    
    // 验证重试次数（初始调用 + 2次重试 = 3次）
    assert_eq!(transport.get_call_count(), 3);
}

#[tokio::test]
async fn test_service_proxy_successful_retry() {
    let transport = Arc::new(MockTransport::new());
    transport.set_should_fail(true);
    
    let registry = Arc::new(MockRegistry::new());
    registry.add_instance("test-service", ServiceInstance::new("test-service", "localhost", 8080, "http"));
    
    let proxy = setup_proxy_with_registry(registry, transport.clone());
    let config = TransportConfig {
        retries: 2,
        retry_delay: 10,
        ..Default::default()
    };
    let request = Request::new("test-service", "GET", "/hello");
    
    // 在第一次调用后修复传输
    tokio::spawn({
        let transport = transport.clone();
        async move {
            tokio::time::sleep(Duration::from_millis(50)).await;
            transport.set_should_fail(false);
        }
    });
    
    let result = proxy.call(request, &config).await;
    // 由于模拟的网络延迟，这里可能仍然失败，但我们主要验证重试逻辑
    // 所以只验证调用次数
    assert!(transport.get_call_count() > 1);
}

#[tokio::test]
async fn test_service_proxy_load_balancing() {
    let transport = Arc::new(MockTransport::new());
    let registry = Arc::new(MockRegistry::new());
    
    // 添加多个实侌
    registry.add_instance("test-service", ServiceInstance::new("test-service", "host1", 8080, "http"));
    registry.add_instance("test-service", ServiceInstance::new("test-service", "host2", 8080, "http"));
    registry.add_instance("test-service", ServiceInstance::new("test-service", "host3", 8080, "http"));
    
    let proxy = setup_proxy_with_registry(registry, transport.clone());
    let config = TransportConfig::default();
    
    // 发起多次调用，验证负载均衡是否生效
    for _ in 0..6 {
        let request = Request::new("test-service", "GET", "/hello");
        let result = proxy.call(request, &config).await;
        assert!(result.is_ok());
    }
    
    assert_eq!(transport.get_call_count(), 6);
}

#[tokio::test]
async fn test_service_proxy_with_unhealthy_instances() {
    let transport = Arc::new(MockTransport::new());
    let registry = Arc::new(MockRegistry::new());
    
    // 添加健康和不健康的实侌
    let mut healthy_instance = ServiceInstance::new("test-service", "healthy-host", 8080, "http");
    let mut unhealthy_instance = ServiceInstance::new("test-service", "unhealthy-host", 8080, "http");
    unhealthy_instance.healthy = false;
    
    registry.add_instance("test-service", healthy_instance);
    registry.add_instance("test-service", unhealthy_instance);
    
    let proxy = setup_proxy_with_registry(registry, transport.clone());
    let config = TransportConfig::default();
    let request = Request::new("test-service", "GET", "/hello");
    
    let result = proxy.call(request, &config).await;
    assert!(result.is_ok());
    
    // 应该只调用健康的实侌
    assert_eq!(transport.get_call_count(), 1);
}

#[tokio::test]
async fn test_service_proxy_different_response_codes() {
    let transport = Arc::new(MockTransport::new());
    let registry = Arc::new(MockRegistry::new());
    registry.add_instance("test-service", ServiceInstance::new("test-service", "localhost", 8080, "http"));
    
    // 测试4xx错误（不应该重试）
    let mut client_error = Response::new(400);
    client_error.set_body(b"Bad Request".to_vec());
    transport.set_response("/client-error", client_error);
    
    let proxy = setup_proxy_with_registry(registry.clone(), transport.clone());
    let config = TransportConfig {
        retries: 2,
        ..Default::default()
    };
    
    let request = Request::new("test-service", "GET", "/client-error");
    let result = proxy.call(request, &config).await;
    
    assert!(result.is_ok());
    let response = result.unwrap();
    assert_eq!(response.status, 400);
    
    // 4xx错误不应该重试，所以只调用一次
    assert_eq!(transport.get_call_count(), 1);
}

#[tokio::test]
async fn test_service_proxy_timeout_behavior() {
    let transport = Arc::new(MockTransport::new());
    let registry = Arc::new(MockRegistry::new());
    registry.add_instance("test-service", ServiceInstance::new("test-service", "localhost", 8080, "http"));
    
    let proxy = setup_proxy_with_registry(registry, transport.clone());
    
    // 设置很短的超时时间和重试参数
    let config = TransportConfig {
        timeout: 100, // 100ms
        retries: 1,
        retry_delay: 50,
        ..Default::default()
    };
    
    let request = Request::new("test-service", "GET", "/hello");
    let result = proxy.call(request, &config).await;
    
    // 应该能成功（mock transport立即返回）
    assert!(result.is_ok());
}

#[tokio::test]
async fn test_service_proxy_batch_calls() {
    let transport = Arc::new(MockTransport::new());
    let registry = Arc::new(MockRegistry::new());
    registry.add_instance("test-service", ServiceInstance::new("test-service", "localhost", 8080, "http"));
    
    let config = TransportConfig::default();
    
    // 并发调用
    let mut handles = Vec::new();
    for i in 0..5 {
        let registry_clone = registry.clone();
        let transport_clone = transport.clone();
        let config_clone = config.clone();
        let handle = tokio::spawn(async move {
            let proxy = setup_proxy_with_registry(registry_clone, transport_clone);
            let request = Request::new("test-service", "GET", &format!("/hello/{}", i));
            proxy.call(request, &config_clone).await
        });
        handles.push(handle);
    }
    
    // 等待所有调用完成
    for handle in handles {
        let result = handle.await.unwrap();
        assert!(result.is_ok());
    }
    
    assert_eq!(transport.get_call_count(), 5);
}